from nnmnkwii.datasets import FileDataSource, FileSourceDataset
from nnmnkwii.datasets import MemoryCacheFramewiseDataset
from nnmnkwii.preprocessing import trim_zeros_frames, remove_zeros_frames
from nnmnkwii.preprocessing import minmax, meanvar, minmax_scale, scale
from nnmnkwii import paramgen
from nnmnkwii.io import hts
from nnmnkwii.frontend import merlin as fe
from nnmnkwii.postfilters import merlin_post_filter

from os.path import join, expanduser, basename, splitext, basename, exists
import os
from glob import glob
import numpy as np
from scipy.io import wavfile
from sklearn.model_selection import train_test_split
import pyworld
import pysptk
import librosa
import librosa.display
import IPython
from IPython.display import Audio
import config as cfg


DATA_ROOT = "./data/slt_arctic_full_data"
max_num_files = None
order = 59
frame_period = 5
windows = [
    (0, 0, np.array([1.0])),
    (1, 1, np.array([-0.5, 0.0, 0.5])),
    (1, 1, np.array([1.0, -2.0, 1.0])),
]


class DurationFeatureSource(FileDataSource):
    def __init__(self, use_phone_alignment=False):
        self.use_phone_alignment = use_phone_alignment

    def collect_files(self):
        if self.use_phone_alignment:
            files = sorted(glob(join(DATA_ROOT, "label_phone_align", "*.lab")))
        else:
            files = sorted(glob(join(DATA_ROOT, "label_state_align", "*.lab")))
        if max_num_files is not None and max_num_files > 0:
            return files[:max_num_files]
        else:
            return files

    def collect_features(self, path):
        # print('path', path)
        labels = hts.load(path)
        # print('labels', labels)
        features = fe.duration_features(labels)
        # print('features', features)
        indices = labels.silence_phone_indices()
        # print('indices', indices)
        features = np.delete(features, indices, axis=0)
        return features.astype(np.float32)


Y_duration_source = DurationFeatureSource(
        use_phone_alignment=False)

Y_duration = FileSourceDataset(Y_duration_source)

for ele in Y_duration:
    print(ele)
    break








for typ in ["acoustic", "duration"]:
    for phase in ["train", "test"]:
        print(typ, phase)
        print(np.max(ld.x[typ][phase][0]))
        print(np.max(ld.y[typ][phase][0]))


for typ in ["acoustic", "duration"]:
    for phase in ["train", "test"]:
        print(typ, phase)
        print(np.min(ld.x[typ][phase][0]))
        print(np.min(ld.y[typ][phase][0]))